from django.core.management.base import BaseCommand
from pymongo import MongoClient
from django.conf import settings
from elasticsearch_dsl import connections
from es.models import TPADocument

class Command(BaseCommand):
    help = 'Sync data from MongoDB to Elasticsearch'

    def handle(self, *args, **kwargs):
        # 连接 MongoDB
        mongo_client = MongoClient(settings.MONGODB_URI)
        db = mongo_client[settings.MONGODB_DB]
        collection = db[settings.MONGODB_COLLECTION]

        # 连接 Elasticsearch
        connections.create_connection(hosts=[settings.ELASTICSEARCH_HOST])

        # 同步数据
        for doc in collection.find({}, {"_id": 1, "problem": 1}):
            es_doc = TPADocument(
                mongo_id=str(doc['_id']),
                problem=doc['problem']
            )
            es_doc.save()

        self.stdout.write(self.style.SUCCESS('Data sync completed'))


# 运行脚本命令 python manage.py sync_data